package matchbox;

import java.util.ArrayList;
import java.util.Random;

import probability.GaussianDistribution;
import utils.Pair;
import basics.DenseVector;
import basics.FeaturesSet;
import basics.Vector;
import basics.VectorMatrix;
import basics.VectorMatrixUtils;
import bayesian.SimpleBayesian;

public class Sampling extends Classifier {

	private static int SAMPLE_SIZE = 1000;
	private static int IGNOE_SAMPLES = 100;

	private static GaussianDistribution delta(GaussianDistribution[] vars, Vector coef) {
		// ArrayList<GaussianDistribution> samples = new ArrayList<GaussianDistribution>();
		//
		// for (int sample = 0; sample < SAMPLE_SIZE; sample++) {
		GaussianDistribution g = new GaussianDistribution(0, 0);
		for (int i = 0; i < vars.length; i++) {
			DenseVector s = vars[i].sample(SAMPLE_SIZE);
			double m = VectorMatrixUtils.mean(s);
			double v = VectorMatrixUtils.stddev(s, m);
			v = v * v;
			double c = coef != null ? coef.get(i) : 1.;
			g.setMu(g.getMu() + m * c);
			g.setSigma2(g.getSigma2() + v * c * c);
		}
		// samples.add(g);
		// }

		return g; // GaussianDistribution.mean(samples, IGNOE_SAMPLES);
	}

	private static GaussianDistribution[] pst_uv(GaussianDistribution[][] uv, Vector xy) {
		GaussianDistribution g[] = new GaussianDistribution[uv.length];
		for (int k = 0; k < uv[0].length; k++) {
			// g[i] = GaussianDistribution.product(uv[i]);
			g[k] = delta(uv[k], xy);
		}
		return g; // delta(g, xy);
	}

	private static GaussianDistribution pz_st(GaussianDistribution s, GaussianDistribution t) {
		double var = (Math.pow(s.getMu(), 2) * t.getSigma2()) + (Math.pow(t.getMu(), 2) * s.getSigma2())
				+ (s.getSigma2() * t.getSigma2());
		return new GaussianDistribution(s.getMu() * t.getMu(), var);
	}

	public static GaussianDistribution pr(double r) {
		return new GaussianDistribution(r, .1);
	}

	private static GaussianDistribution pst_z(GaussianDistribution st_other, GaussianDistribution z) {
		return GaussianDistribution.sampleIndicatorDivide(5, z, st_other);
	}

	private static GaussianDistribution[][] puv_st(GaussianDistribution[][] uv, GaussianDistribution[] st, Vector xy) {
		ArrayList<ArrayList<ArrayList<GaussianDistribution>>> uvs = new ArrayList<ArrayList<ArrayList<GaussianDistribution>>>();

		for (int i = 0; i < uv.length; i++) {
			uvs.add(new ArrayList<ArrayList<GaussianDistribution>>());
			for (int k = 0; k < uv[0].length; k++) {
				uvs.get(i).add(new ArrayList<GaussianDistribution>());
			}
		}
		for (int i = 0; i < uv.length; i++) {
			if (xy.get(i) == 0)
				continue;

			GaussianDistribution[] uvk = new GaussianDistribution[uv[0].length];
			for (int k = 0; k < uv[0].length; k++) {
				uvk[k] = uv[i][k];
			}

			for (int k = 0; k < uv[0].length; k++) {
				GaussianDistribution sum = delta(uvk, xy);
				sum = sum.minus(uv[i][k]);
				sum = st[k].add(sum);
				uv[i][k] = sum.times(1. / xy.get(i));
				uvs.get(i).get(k).add(uv[i][k]);
			}
		}

		for (int i = 0; i < uv.length; i++) {
			for (int k = 0; k < uv[0].length; k++) {
				uv[i][k] = GaussianDistribution.mean(uvs.get(i).get(k), IGNOE_SAMPLES);
			}
		}
		return uv;
	}

	private GaussianDistribution[] pz_r(GaussianDistribution pr) {
		ArrayList<ArrayList<GaussianDistribution>> zs = new ArrayList<ArrayList<GaussianDistribution>>();
		GaussianDistribution sum = GaussianDistribution.sum(_z);
		for (int k = 0; k < _k; k++) {
			zs.add(new ArrayList<GaussianDistribution>());
		}

		for (int i = 0; i < SAMPLE_SIZE; i++) {
			for (int k = 0; k < _k; k++) {
				_z[k] = pr.minus(sum.minus(_z[k]));
				zs.get(k).add(_z[k]);
			}
		}

		for (int k = 0; k < _k; k++) {
			_z[k] = GaussianDistribution.mean(zs.get(k), IGNOE_SAMPLES);
		}

		return _z;
	}

	private GaussianDistribution pr_z(GaussianDistribution[] z) {
		GaussianDistribution r = delta(z, null);
		if (r.getMu() > 5)
			r.setMu(5.);
		else if (r.getMu() < 1)
			r.setMu(1);

		return r;
	}

	private GaussianDistribution forward(Vector x, Vector y) {
		_z = new GaussianDistribution[_k];
		_s = pst_uv(_U, x);
		_t = pst_uv(_V, y);

		for (int k = 0; k < _k; k++) {
			_z[k] = pz_st(_s[k], _t[k]);
		}
		return pr_z(_z);
	}

	private void backward(Vector x, Vector y, double r) {
		_z = pz_r(pr(r));

		for (int k = 0; k < _k; k++) {
			ArrayList<GaussianDistribution> ss = new ArrayList<GaussianDistribution>();
			ArrayList<GaussianDistribution> st = new ArrayList<GaussianDistribution>();
			for (int i = 0; i < SAMPLE_SIZE; i++) {
				_t[k] = pst_z(_s[k], _z[k]);
				st.add(_t[k]);
				_s[k] = pst_z(_t[k], _z[k]);
				ss.add(_s[k]);
			}
			_s[k] = GaussianDistribution.mean(ss, IGNOE_SAMPLES);
			_t[k] = GaussianDistribution.mean(st, IGNOE_SAMPLES);
		}

		puv_st(_U, _s, x);
		puv_st(_V, _t, y);
	}

	GaussianDistribution _z[];
	GaussianDistribution _s[];
	GaussianDistribution _t[];
	private GaussianDistribution[][] _U, _V;
	private int _k;

	private GaussianDistribution[][] init(int n, int m) {
		GaussianDistribution[][] p = new GaussianDistribution[n][m];
		Random r = new Random();
		for (int i = 0; i < n; i++) {
			for (int j = 0; j < m; j++) {
				p[i][j] = new GaussianDistribution(r.nextDouble(), 10.);
			}
		}
		return p;
	}

	public void train(FeaturesSet X, FeaturesSet Y, FeaturesSet R, int lDimension) {
		_U = init(X.get(0).length(), lDimension);
		_V = init(Y.get(0).length(), lDimension);
		_k = lDimension;

		for (int i = 0; i < X.size(); i++) {
			for (int j = 0; j < Y.size(); j++) {
				if (R.get(i).get(j) != 0) {
					System.out.println(i + "," + j);
					Vector x = (Vector) X.get(i);
					Vector y = (Vector) Y.get(j);
					double r = R.get(i).get(j);
					GaussianDistribution pr = null, pr_prev;
					do {
						pr_prev = pr;
						pr = forward(x, y);
						backward(x, y, r);
					} while (pr_prev == null || pr.getMu() - pr_prev.getMu() > .5);
				}
			}
		}
	}

	@Override
	public double infer(VectorMatrix u, VectorMatrix it) {
		return forward((Vector) u, (Vector) it).getMu();
	}

	@Override
	protected void trainUserItem(VectorMatrix u, VectorMatrix it, double r) {
		forward((Vector) u, (Vector) it);
		backward((Vector) u, (Vector) it, r);
	}

}
